#include "StdAfx.h"
#include "PoseModifierHelper.h"

namespace PoseModifierHelper {


	int32 IK_Solver( const CModelSkeleton* __restrict pModelSkeleton, const char* strLimb, const QuatT& vLocalPos, QuatT* const __restrict pRelPose,	QuatT* const __restrict	pAbsPose, const bool adjustEndEffector )
	{
		uint64 nHandle = *(uint64*)strLimb;
		int32 idxDefinition = pModelSkeleton->GetLimbDefinitionIdx(nHandle);
		if (idxDefinition < 0)
		{
			//if (Console::GetInst().ca_AnimWarningLevel > 2)
			//{
			//	CryLogAlways("IK Solver: Limb Definition %s does not exist. Aborting.", SPU_MAIN_PTR(strLimb));
			//}
			return 0;
		}

		const IKLimbType& rIKLimbType = pModelSkeleton->m_IKLimbTypes[idxDefinition];
		if (rIKLimbType.m_nSolver == *(uint32*)SPU_MAIN_PTR("2BIK"))
			IK_Solver2Bones(vLocalPos.t, rIKLimbType, pRelPose, pAbsPose);
		if (rIKLimbType.m_nSolver == *(uint32*)SPU_MAIN_PTR("3BIK"))
			IK_Solver3Bones(vLocalPos.t, rIKLimbType, pRelPose, pAbsPose);
		if (rIKLimbType.m_nSolver == *(uint32*)SPU_MAIN_PTR("CCDX"))
			IK_SolverCCD(vLocalPos.t, rIKLimbType, pRelPose, pAbsPose);

		const CModelJoint* __restrict parrModelJoints = &pModelSkeleton->m_arrModelJoints[0];
		uint32 numChilds = rIKLimbType.m_arrLimbChildren.size();
		const int16* pJointIdx = &rIKLimbType.m_arrLimbChildren[0];

		if (adjustEndEffector)
		{
			uint32 numLinks=rIKLimbType.m_arrJointChain.size();
			int32 nEndEffector =	rIKLimbType.m_arrJointChain[numLinks-1].m_idxJoint;
			int32 nParentEndEffector =	parrModelJoints[nEndEffector].m_idxParent;
			pAbsPose[nEndEffector].q = vLocalPos.q;
			pRelPose[nEndEffector].q = !pAbsPose[nParentEndEffector].q * pAbsPose[nEndEffector].q;
		}

		for (uint32 i=0; i<numChilds; ++i)
		{
			int32 c = pJointIdx[i];
			int32 p = parrModelJoints[c].m_idxParent;
			pAbsPose[c] = pAbsPose[p] * pRelPose[c];
			ANIM_ASSET_ASSERT( pAbsPose[c].q.IsUnit() );
		}

#ifdef _DEBUG
		uint32 numJoints=pModelSkeleton->GetJointCount();
		for (uint32 j=0; j<numJoints; j++)
		{
			ANIM_ASSET_ASSERT( pRelPose[j].q.IsUnit() );
			ANIM_ASSET_ASSERT( pAbsPose[j].q.IsUnit() );
			ANIM_ASSET_ASSERT( pRelPose[j].IsValid() );
			ANIM_ASSET_ASSERT( pAbsPose[j].IsValid() );
		}
#endif			

		return 1;
	};

void IK_Solver2Bones( const Vec3r& goal, const IKLimbType& rIKLimbType,QuatT* const __restrict	pRelPose,QuatT* const __restrict	pAbsPose)
{
	int32 b0 = rIKLimbType.m_arrJointChain[0].m_idxJoint; assert(b0>0); //BaseRoot
	int32 b1 = rIKLimbType.m_arrJointChain[1].m_idxJoint; assert(b1>0); //Root of Chain
	int32 b2 = rIKLimbType.m_arrJointChain[2].m_idxJoint; assert(b2>0); //Middle Joint 
	int32 b3 = rIKLimbType.m_arrJointChain[3].m_idxJoint; assert(b3>0); //EndEffector

	uint32 numLinks = rIKLimbType.m_arrRootToEndEffector.size();
	pAbsPose[0]	= pRelPose[0];
	ANIM_ASSET_ASSERT( pRelPose[0].q.IsUnit() );
	for (uint32 i=1; i<numLinks; i++)
	{
		int32 p=rIKLimbType.m_arrRootToEndEffector[i-1];
		int32 c=rIKLimbType.m_arrRootToEndEffector[i];
		ANIM_ASSET_ASSERT( pRelPose[c].q.IsUnit() );
		ANIM_ASSET_ASSERT( pRelPose[p].q.IsUnit() );
		pAbsPose[c]	= pAbsPose[p] * pRelPose[c];
		ANIM_ASSET_ASSERT( pRelPose[c].q.IsUnit() );
	}

	Vec3r t0 = pAbsPose[b1].t;
	Vec3r t1 = pAbsPose[b2].t;
	Vec3r t2 = pAbsPose[b3].t;
	if ( ((goal-t2)|(goal-t2)) < real(0.00000001) )
		return; //end-effector and new goal are very close

	Vec3r h,axis;
	Vec3r a = t1-t0;
	Vec3r b = t2-t1;
	Vec3r c = goal-t0;
	Vec3r cross=a%b;
	if ((cross|cross) < real(0.00001) )
		return;

	cross.Normalize();
	ANIM_ASSET_ASSERT((fabs_tpl(1-(cross|cross))) < real(0.00001) );   //check if unit-vector

	real daxis;
	real alen = a.GetLength();
	if (alen < real(0.1) )
		return;  //check for invalid configurations / broken bones
	real blen = b.GetLength();
	if (blen < real(0.1) )
		return;  //check for invalid configurations / broken bones
	real clen = c.GetLength();
	if (clen < real(0.01) )
		return; //invalid target

	real cos0	= clamp(	(alen*alen+clen*clen-blen*blen)/(real(2.0)*alen*clen) ,real(-1.0),real(+1.0) );
	real sin0	=	sqrt( max(real(1.0)-cos0*cos0,real(0.0)) );

	a/=alen;  //normalize
	a = a.GetRotated( cross, cos0,sin0 );
	ANIM_ASSET_ASSERT((fabs_tpl(1-(a|a))) < real(0.00001) );   //check if unit-vector

	pAbsPose[b0].q.NormalizeSafe();

	c/=clen;  //normalize
	ANIM_ASSET_ASSERT((fabs_tpl(1-(c|c))) < real(0.00001) );   //check if unit-vector
	h=(a+c).GetNormalized();
	ANIM_ASSET_ASSERT((fabs_tpl(1-(h|h))) < real(0.00001) );   //check if unit-vector

	axis=a%c;	daxis=axis|axis;
	if (daxis>real(0.000001))
		pAbsPose[b1].q = Quat(f32(a|h), axis/sqrt(daxis)*(a%h).GetLength()) * pAbsPose[b1].q;
	ANIM_ASSET_ASSERT( pAbsPose[b1].q.IsUnit() );

	pAbsPose[b1].q.NormalizeSafe();
	pRelPose[b1].q = !pAbsPose[b0].q * pAbsPose[b1].q;
	ANIM_ASSET_ASSERT( pRelPose[b1].q.IsUnit() );
	pAbsPose[b2]  = pAbsPose[b1] * pRelPose[b2];
	ANIM_ASSET_ASSERT( pAbsPose[b2].q.IsUnit() );
	pAbsPose[b3]  = pAbsPose[b2] * pRelPose[b3];
	ANIM_ASSET_ASSERT( pAbsPose[b3].q.IsUnit() );

	//---------------------------------------------------

	c = (goal-pAbsPose[b2].t).GetNormalized(); 
	ANIM_ASSET_ASSERT((fabs_tpl(1-(c|c))) < real(0.00001) );   //check if unit-vector
	b = Vec3r(pAbsPose[b3].t-pAbsPose[b2].t)/blen;
	h = (b+c).GetNormalized();
	ANIM_ASSET_ASSERT((fabs_tpl(1-(h|h))) < real(0.00001) );   //check if unit-vector

	axis=b%c;	daxis=axis|axis;
	if (daxis>real(0.000001))
		pAbsPose[b2].q = Quat(f32(b|h),axis/sqrt(daxis)*(b%h).GetLength())*pAbsPose[b2].q;

	ANIM_ASSET_ASSERT( pAbsPose[b2].q.IsUnit() );
	pAbsPose[b2].q.NormalizeSafe();
	pRelPose[b2].q = !pAbsPose[b1].q * pAbsPose[b2].q;
	ANIM_ASSET_ASSERT( pRelPose[b2].q.IsUnit() );
}

void IK_Solver3Bones(const Vec3r& goal, const IKLimbType& rIKLimbType,QuatT* const __restrict pRelPose,QuatT* const __restrict pAbsPose)
{
	int32 b0 = rIKLimbType.m_arrJointChain[0].m_idxJoint; ANIM_ASSET_ASSERT(b0>0); //BaseRoot
	int32 b1 = rIKLimbType.m_arrJointChain[1].m_idxJoint; ANIM_ASSET_ASSERT(b1>0); //Root of Chain
	int32 b2 = rIKLimbType.m_arrJointChain[2].m_idxJoint; ANIM_ASSET_ASSERT(b2>0); //Middle Joint 
	int32 b3 = rIKLimbType.m_arrJointChain[3].m_idxJoint; ANIM_ASSET_ASSERT(b3>0); //Middle Joint 
	int32 b4 = rIKLimbType.m_arrJointChain[4].m_idxJoint; ANIM_ASSET_ASSERT(b4>0); //EndEffector

	uint32 numLinks = rIKLimbType.m_arrRootToEndEffector.size();
	pAbsPose[0]	= pRelPose[0];
	for (uint32 i=1; i<numLinks; i++)
	{
		int32 p=rIKLimbType.m_arrRootToEndEffector[i-1];
		int32 c=rIKLimbType.m_arrRootToEndEffector[i];
		ANIM_ASSET_ASSERT( pRelPose[c].q.IsUnit() );
		ANIM_ASSET_ASSERT( pAbsPose[p].q.IsUnit() );
		pAbsPose[c]	= pAbsPose[p] * pRelPose[c];
		ANIM_ASSET_ASSERT( pAbsPose[c].q.IsUnit() );
	}

	Vec3r t1 = pAbsPose[b1].t;
	Vec3r t2 = pAbsPose[b2].t;
	Vec3r t4 = pAbsPose[b4].t;
	if ( ((goal-t4)|(goal-t4)) < real(0.0000001) )
		return; //end-effector and new goal are very close

	Vec3r h,axis; real daxis;
	Vec3r a=t2-t1;		real alen=a.GetLength();
	if (alen < real(0.1) )
		return;  //check for invalid configurations / broken bones
	Vec3r b=t4-t2;		real blen=b.GetLength();
	if (blen < real(0.1) )
		return;  //check for invalid configurations / broken bones
	Vec3r c=goal-t1;	real clen=c.GetLength();
	if (clen < real(0.01) )
		return;  //invalid target

	f64 ration = (t1-t4).GetLength()/clen-real(1.0);
	if (ration<0)	ration*=3.5f;  //open

	Vec3r zaxis = pRelPose[b3].q.GetColumn2();
	Quat zrot; zrot.SetRotationAA(clamp(f32(ration),-1.1f,+0.4f),zaxis); 	
	pRelPose[b3].q = pRelPose[b3].q * zrot;
	ANIM_ASSET_ASSERT( pRelPose[b3].q.IsUnit() );
	pAbsPose[b3]  = pAbsPose[b2] * pRelPose[b3];
	ANIM_ASSET_ASSERT( pAbsPose[b3].q.IsUnit() );
	pAbsPose[b4]  = pAbsPose[b3] * pRelPose[b4];
	ANIM_ASSET_ASSERT( pAbsPose[b4].q.IsUnit() );

	t1 = pAbsPose[b1].t;
	t2 = pAbsPose[b2].t;
	t4 = pAbsPose[b4].t;
	a=t2-t1;		alen=a.GetLength();
	b=t4-t2;		blen=b.GetLength();
	c=goal-t1;	clen=c.GetLength();
	Vec3r cross=a%b;
	if ((cross|cross) < real(0.00001) )
		return;
	cross.Normalize();
	ANIM_ASSET_ASSERT((fabs_tpl(1-(cross|cross))) < real(0.00001) );   //check if unit-vector

	real cos0	= clamp(	(real)((alen*alen+clen*clen-blen*blen)/(real(2.0)*alen*clen)) ,real(-1.0),real(+1.0) );
	real sin0	=	sqrt( max(real(1.0)-cos0*cos0,real(0.0)) );

	a/=alen;
	a = a.GetRotated( cross, cos0,sin0 );
	ANIM_ASSET_ASSERT((fabs_tpl(1-(a|a))) < real(0.00001) );   //check if unit-vector

	pAbsPose[b0].q.NormalizeSafe();

	c/=clen;
	ANIM_ASSET_ASSERT((fabs_tpl(1-(c|c))) < real(0.00001) );   //check if unit-vector
	h=(a+c).GetNormalized();
	ANIM_ASSET_ASSERT((fabs_tpl(1-(h|h))) < real(0.00001) );   //check if unit-vector
	axis=a%c; daxis=axis|axis;
	if (daxis>real(0.00001))
		pAbsPose[b1].q = Quat(f32(a|h),axis/sqrt(daxis)*(a%h).GetLength()) * pAbsPose[b1].q;
	pAbsPose[b1].q.NormalizeSafe();
	ANIM_ASSET_ASSERT( pAbsPose[b1].q.IsUnit() );
	pRelPose[b1].q = !pAbsPose[b0].q * pAbsPose[b1].q;
	ANIM_ASSET_ASSERT( pRelPose[b1].q.IsUnit() );
	pAbsPose[b2]  = pAbsPose[b1] * pRelPose[b2];
	ANIM_ASSET_ASSERT( pAbsPose[b2].q.IsUnit() );
	pAbsPose[b3]  = pAbsPose[b2] * pRelPose[b3];
	ANIM_ASSET_ASSERT( pAbsPose[b3].q.IsUnit() );
	pAbsPose[b4]  = pAbsPose[b3] * pRelPose[b4];
	ANIM_ASSET_ASSERT( pAbsPose[b4].q.IsUnit() );

	//---------------------------------------------------

	c = (goal-pAbsPose[b2].t).GetNormalized(); 
	ANIM_ASSET_ASSERT((fabs_tpl(1-(c|c))) < real(0.00001) );   //check if unit-vector
	b = Vec3r(pAbsPose[b4].t-pAbsPose[b2].t)/blen;
	h = (b+c).GetNormalized();
	ANIM_ASSET_ASSERT((fabs_tpl(1-(h|h))) < real(0.00001) );   //check if unit-vector

	axis=b%c;	daxis=axis|axis;
	if (daxis>real(0.00001))
		pAbsPose[b2].q = Quat(f32(b|h),axis/sqrt(daxis)*(b%h).GetLength())*pAbsPose[b2].q;
	ANIM_ASSET_ASSERT( pAbsPose[b2].q.IsUnit() );
	pAbsPose[b2].q.NormalizeSafe();
	pRelPose[b2].q = !pAbsPose[b1].q * pAbsPose[b2].q;
	ANIM_ASSET_ASSERT( pRelPose[b2].q.IsUnit() );
}

//----------------------------------------------------------------------------------------------
//----------------------------------------------------------------------------------------------
//----------------------------------------------------------------------------------------------
void IK_SolverCCD(const Vec3& vTarget, const IKLimbType& rIKLimbType,QuatT* const __restrict pRelPose,QuatT* const __restrict	pAbsPose)
{
	uint32 numR2E = rIKLimbType.m_arrRootToEndEffector.size();
	pAbsPose[0]	= pRelPose[0];
	for (uint32 i=1; i<numR2E; i++)
	{
		int32 p=rIKLimbType.m_arrRootToEndEffector[i-1];
		int32 c=rIKLimbType.m_arrRootToEndEffector[i];
		ANIM_ASSET_ASSERT( pRelPose[c].q.IsUnit() );
		ANIM_ASSET_ASSERT( pAbsPose[p].q.IsUnit() );
		pAbsPose[c]	= pAbsPose[p] * pRelPose[c];
		ANIM_ASSET_ASSERT( pAbsPose[c].q.IsUnit() );
	}

	f32 fTransCompensation=0;
	uint32 numLinks		= rIKLimbType.m_arrJointChain.size();
	for (uint32 i=2; i<numLinks; i++)
	{
		int32 p = rIKLimbType.m_arrJointChain[i-1].m_idxJoint;
		int32 c = rIKLimbType.m_arrJointChain[i].m_idxJoint;
		fTransCompensation += (pAbsPose[c].t-pAbsPose[p].t).GetLengthFast();
	}

	f32 inumLinks			= 1.0f/f32(numLinks);
	int32 nRootIdx		=	rIKLimbType.m_arrJointChain[1].m_idxJoint;  //Root
	int32 nEndEffIdx	=	rIKLimbType.m_arrJointChain[numLinks-1].m_idxJoint; //EndEffector
	ANIM_ASSET_ASSERT(nRootIdx<nEndEffIdx);
	int32	iJointIterator	= 1;  //numLinks-2;

	// Cyclic Coordinate Descent
	for (uint32 i=0; i<rIKLimbType.m_nInterations; i++)
	{
		Vec3	vecEnd				=	pAbsPose[nEndEffIdx].t;			// Position of end effector
		int32 parentLinkIdx	= rIKLimbType.m_arrJointChain[iJointIterator-1].m_idxJoint; ANIM_ASSET_ASSERT(parentLinkIdx>=0);
		int32 LinkIdx				= rIKLimbType.m_arrJointChain[iJointIterator].m_idxJoint;
		Vec3	vecLink				=	pAbsPose[LinkIdx].t;	// Position of current node
		Vec3	vecLinkTarget	= (vTarget-vecLink).GetNormalized();					// Vector current link -> target
		Vec3	vecLinkEnd		= (vecEnd-vecLink).GetNormalized();								// Vector current link -> current effector position

		Quat qrel; qrel.SetRotationV0V1(vecLinkEnd,vecLinkTarget); 
		ANIM_ASSET_ASSERT((fabs_tpl(1-(qrel|qrel)))<0.001);   //check if unit-quaternion
		if (qrel.w<0) qrel.w = -qrel.w; 
		
		f32 ji = iJointIterator*inumLinks+0.3f;
		f32 t	=	min(rIKLimbType.m_fStepSize*ji,0.4f);
		qrel.w*=t+1.0f-t; qrel.v*=t;	qrel.NormalizeSafe();

		//calculate new relative IK-orientation 
		pRelPose[LinkIdx].q = !pAbsPose[parentLinkIdx].q * qrel * pAbsPose[LinkIdx].q;
		pRelPose[LinkIdx].q.NormalizeSafe();

		//calculate new absolute IK-orientation 
		for (uint32 j=iJointIterator; j<numLinks; j++)
		{
			int32 c = rIKLimbType.m_arrJointChain[j].m_idxJoint;
			int32 p = rIKLimbType.m_arrJointChain[j-1].m_idxJoint;	assert(p>=0);
			ANIM_ASSET_ASSERT( pRelPose[c].q.IsUnit() );
			ANIM_ASSET_ASSERT( pAbsPose[p].q.IsUnit() );
			pAbsPose[c]= pAbsPose[p]*pRelPose[c];
			ANIM_ASSET_ASSERT( pAbsPose[c].q.IsUnit() );
		}

		f32	fError	= (pAbsPose[nEndEffIdx].t-vTarget).GetLength();
		/*
		static Ang3 angle=Ang3(ZERO);
		angle.x += 0.1f;
		angle.y += 0.01f;
		angle.z += 0.001f;
		AABB sAABB = AABB(Vec3(-0.1f,-0.1f,-0.1f),Vec3(+0.1f,+0.1f,+0.1f));
		OBB obb =	OBB::CreateOBBfromAABB( Matrix33::CreateRotationXYZ(angle), sAABB );
		g_pAuxGeom->DrawOBB(obb,pAbsPose[nEndEffIdx].t,1,RGBA8(0xff,0x00,0x00,0xff),eBBD_Extremes_Color_Encoded);
		*/

		//	float fColor[4] = {0,1,0,1};
		//	g_pIRenderer->Draw2dLabel( 1,g_YLine, 1.2f, fColor, false,"nUpdateSteps: %d   fError: %f  iJointIterator: %d  ji: %f  t: %f",i,fError,iJointIterator,ji,t );	
		//	g_YLine+=12.0f;

		if (fError < rIKLimbType.m_fThreshold) break; 

		iJointIterator++;	//Next link
		if( iJointIterator > (int)(numLinks)-3)				// is it correct, uint - 3 can cause sirious problems
			iJointIterator = 1;
	}

	//	float fColor[4] = {0,1,0,1};
	//	f32	fError1	= (pAbsPose[nEndEffIdx].t-vTarget).GetLength();		
	//	g_pIRenderer->Draw2dLabel( 1,g_YLine, 1.4f, fColor, false,"fError1: %f",fError1);	
	//	g_YLine+=12.0f;

	//-----------------------------------------------------------------------------
	//--do a cheap translation compensation to fix the error
	//-----------------------------------------------------------------------------
	Vec3 absEndEffector =	pAbsPose[nEndEffIdx].t;			// Position of end effector
	Vec3 vDistance	= (vTarget-absEndEffector); //f32(numCCDJoints);
	f32 fDistance = vDistance.GetLengthFast();
	if (fDistance>fTransCompensation)
		vDistance*=fTransCompensation/fDistance;
	Vec3 bPartDistance	=	vDistance/f32(numLinks-0);
	Vec3 vAddDistance =	bPartDistance;

	for (uint32 i=1; i<numLinks; i++)
	{
		int c = rIKLimbType.m_arrJointChain[i].m_idxJoint;
		int p = rIKLimbType.m_arrJointChain[i-1].m_idxJoint; assert(p>=0);
		pAbsPose[c].t += vAddDistance;
		vAddDistance += bPartDistance;
		ANIM_ASSET_ASSERT( pAbsPose[c].q.IsUnit() );
		ANIM_ASSET_ASSERT( pAbsPose[p].q.IsUnit() );
		pRelPose[c] = pAbsPose[p].GetInverted() * pAbsPose[c];
		pRelPose[c].q.NormalizeSafe();
	}

	//	f32	fError2	= (pAbsPose[nEndEffIdx].t-vTarget).GetLength();	
	//	g_pIRenderer->Draw2dLabel( 1,g_YLine, 1.4f, fColor, false,"fError2: %f",fError2);	
	//	g_YLine+=12.0f;
}


} // namespace PoseModifierHelper
